package com.chj.zhongji.class09.test01;

import java.util.HashSet;

//https://blog.csdn.net/u012017783/article/details/80049539
public class Code04_TopKSumCrossTwoArrays {
	public static class HeapNode {
		public int row;
		public int col;
		public int value;

		public HeapNode(int row, int col, int value) {
			this.row = row;
			this.col = col;
			this.value = value;
		}
	}

	public static int[] topKSum(int[] a1, int[] a2, int topK) {
		if (a1 == null || a2 == null || topK < 1) {
			return null;
		}

		topK = Math.min(topK, a1.length * a2.length);
		HeapNode[] heap = new HeapNode[topK + 1];
		int heapSize = 0;
		int headR = a1.length - 1;
		int headC = a2.length - 1;
		int uR = -1;
		int uC = -1;
		int lR = -1;
		int lC = -1;
		heapInsert(heap, heapSize, headR, headC, a1[headR] + a2[headC]);
		heapSize++;
//		HashSet<String> positionSet = new HashSet<String>();
		boolean[][] set = new boolean[a1.length][a2.length];
		int[] res = new int[topK];
		int resIndex = 0;
		while (resIndex != topK) {
			HeapNode head = popHead(heap, heapSize);
			heapSize--;
			res[resIndex++] = head.value;

			headR = head.row;
			headC = head.col;

			uR = headR - 1;
			uC = headC;
			if (headR != 0 && !set[uR][uC]) {
				heapInsert(heap, heapSize, uR, uC, a1[uR] + a2[uC]);
				heapSize++;
				set[uR][uC] = true;
			}

//			if (headR != 0 && !isContains(uR, uC, positionSet)) {
//				heapInsert(heap, heapSize, uR, uC, a1[uR] + a2[uC]);
//				heapSize++;
//				addPositionSet(uR, uC, positionSet);
//			}

			lR = headR;
			lC = headC - 1;
			if (headC != 0 && !set[lR][lC]) {
				heapInsert(heap, heapSize, lR, lC, a1[lR] + a2[lC]);
				heapSize++;
				set[lR][lC] = true;
			}

//			if (headC != 0 && !isContains(lR, lC, positionSet)) {
//				heapInsert(heap, heapSize, lR, lC, a1[lR] + a2[lC]);
//				heapSize++;
//				addPositionSet(lR, lC, positionSet);
//			}
		}
		return res;
	}

	public static HeapNode popHead(HeapNode[] heap, int heapSize) {
		HeapNode res = heap[0];
		swap(heap, 0, heapSize - 1);
		heapSize--;
		heap[heapSize] = null;
		heapify(heap, 0, heapSize);
		return res;
	}

	public static void heapInsert(HeapNode[] heap, int index, int row, int col, int value) {
		heap[index] = new HeapNode(row, col, value);
		int parent = (index - 1) / 2;
		while (index != 0) {
			if (heap[index].value > heap[parent].value) {
				swap(heap, index, parent);
				index = parent;
				parent = (index - 1) / 2;
			} else {
				break;
			}
		}
	}

	public static void heapify(HeapNode[] heap, int index, int heapSize) {
		int left = index * 2 + 1;
		int right = index * 2 + 2;
		int largest = index;
		while (left < heapSize) {
			if (heap[left].value > heap[largest].value) {
				largest = left;
			}

			if (right < heapSize && heap[right].value > heap[largest].value) {
				largest = right;
			}

			if (largest != index) {
				swap(heap, largest, index);
			} else {
				break;
			}
			index = largest;
			left = index * 2 + 1;
			right = index * 2 + 1;
		}
	}

	public static void swap(HeapNode[] heap, int index1, int index2) {
		HeapNode tmp = heap[index1];
		heap[index1] = heap[index2];
		heap[index2] = tmp;
	}

	public static boolean isContains(int row, int col, HashSet<String> set) {
		return set.contains(String.valueOf(row + "_" + col));
	}

	public static void addPositionSet(int row, int col, HashSet<String> set) {
		set.add(String.valueOf(row + "_" + col));
	}

	public static void main(String[] args) {
		int[] arr1 = { 1, 2, 3, 4, 5 };
		int[] arr22 = { 3, 5, 7, 9, 11 };
		int k = 4;
		int[] res = topKSum(arr1, arr22, k);
		for (int i = 0; i < res.length; i++) {
			System.out.print(res[i] + " ");
		}
	}
}
